Skip to content

feat: PyTorch TensorDict Phase 4 - Schema Inference, IO Serialization, and Hypothesis Strategies - #2316

Open
cosmicBboy wants to merge 8 commits into
dev/pytorch-tensordict-phase1from
pr/pytorch-tensordict-phase3-4
Open

feat: PyTorch TensorDict Phase 4 - Schema Inference, IO Serialization, and Hypothesis Strategies#2316
cosmicBboy wants to merge 8 commits into
dev/pytorch-tensordict-phase1from
pr/pytorch-tensordict-phase3-4

Conversation

@cosmicBboy

Copy link
Copy Markdown
Collaborator

PyTorch TensorDict Phase 4 Implementation

This PR implements Phase 4 features for PyTorch TensorDict integration in pandera.

What's New

🎯 Schema Inference

  • Automatically infer schemas from TensorDict/tensorclass data using
  • Detects dtypes, shapes, and value range statistics automatically
  • Supports both regular TensorDict and tensorclass objects

💾 IO Serialization/Deserialization

  • Save/load schemas with YAML ( / )
  • JSON serialization support ( / )
  • Save TensorDicts with embedded schema metadata for data integrity validation

🧪 Hypothesis Strategies

  • Generate synthetic TensorDicts from schema definitions
  • Property-based testing integration with Hypothesis
  • Supports both TensorDict and tensorclass generation

Files Changed

New Modules:

    • Schema inference functionality
    • IO serialization/deserialization
    • Hypothesis strategies

Updated Entry Points:

Tests Added:

  • (12 tests)
  • (9 tests)
  • (5 tests)

Documentation:

  • Added tensordict_schema_inference.md
  • Added tensordict_io.md
  • Added tensordict_strategies.md
  • Updated index.md and existing docs

Testing

All 84 tensordict tests pass:

Related Work**

Based on: (Phase 3 already merged)

Niels Bantilan added 4 commits April 25, 2026 18:14
…ensorDict Phase 3

- Add coerce_dtype implementation in TensorDictSchemaBackend
- Support both schema-level (schema.coerce=True) and tensor-level coerce flags
- Handle both TensorDict and tensorclass objects during coercion
- Added 9 comprehensive tests for coercion functionality
- Updated documentation with type coercion examples

Phase 2 backend infrastructure was already implemented in Phase 1.
Phase 3 adds dtype coercion capability to match other pandera backends.
…ion, and hypothesis strategies

- Implement pa.infer_schema() to automatically infer TensorDict schemas from data
- Add schema.save/load() methods with YAML support for configuration management
- Integrate Hypothesis strategies for property-based testing of TensorDict schemas
- Include complete ML pipeline example demonstrating all Phase 4 features
- Add schema inference via pa.infer_schema() to automatically infer TensorDict schemas from data
- Implement IO operations (save/load/to_yaml/from_yaml) for TensorDictSchema serialization
- Integrate Hypothesis strategies (tensordict_strategy, tensorclass_strategy) for property-based testing
- Update API exports and entry points in pandera/tensordict.py
- Add 30 comprehensive tests covering all Phase 4 features
…trategies)

- Add tensordict_schema_inference.md with pa.infer_schema() examples
- Add tensordict_io.md with YAML/JSON serialization examples
- Add tensordict_strategies.md with Hypothesis property-based testing docs
- Update index.md to include new doc pages
- Fix existing docs to use correct batch_size values
@codecov

codecov Bot commented Apr 28, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 13.57702% with 331 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.83%. Comparing base (8550274) to head (c92ad6a).

Files with missing lines Patch % Lines
pandera/io/tensordict_io.py 15.15% 140 Missing ⚠️
pandera/schema_inference/tensordict.py 0.00% 72 Missing ⚠️
pandera/strategies/tensordict_strategies.py 5.88% 64 Missing ⚠️
pandera/backends/tensordict/base.py 0.00% 30 Missing ⚠️
pandera/tensordict.py 0.00% 16 Missing ⚠️
pandera/strategies/__init__.py 60.00% 8 Missing ⚠️
pandera/schema_inference/__init__.py 80.00% 1 Missing ⚠️
Additional details and impacted files
@@                        Coverage Diff                        @@
##           dev/pytorch-tensordict-phase1    #2316      +/-   ##
=================================================================
- Coverage                          78.39%   77.83%   -0.56%     
=================================================================
  Files                                193      199       +6     
  Lines                              15925    16297     +372     
=================================================================
+ Hits                               12484    12685     +201     
- Misses                              3441     3612     +171     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

- Add schema inference via pa.infer_schema() to automatically infer TensorDict schemas from data
- Implement IO operations (save/load/to_yaml/from_yaml) for TensorDictSchema serialization
- Integrate Hypothesis strategies (tensordict_strategy, tensorclass_strategy) for property-based testing
- Update API exports and entry points in pandera/tensordict.py
- Add 30 comprehensive tests covering all Phase 4 features

Docs:
- Add tensordict_schema_inference.md with pa.infer_schema() examples
- Add tensordict_io.md with YAML/JSON serialization examples
- Add tensordict_strategies.md with Hypothesis property-based testing docs
- Update index.md to include new doc pages
@cosmicBboy
cosmicBboy force-pushed the pr/pytorch-tensordict-phase3-4 branch from b684bfe to c2a33ea Compare April 28, 2026 01:20
Niels Bantilan added 3 commits April 27, 2026 22:38
- Add proper type annotations to _strategy_import_error decorator
- Fix statistics dict typing in io/tensordict_io.py
- Resolve single overload issue in schema_inference/tensordict.py
- Handle optional infer_schema type in entry point
…dict-phase4

# Conflicts:
#	docs/source/pytorch_guide/tensordict_io.md
#	docs/source/pytorch_guide/tensordict_schema.md
#	docs/source/pytorch_guide/tensordict_strategies.md
- Fix _wrapper decorator type annotations in strategies module
- Remove redundant overload signatures in schema inference
- Add TYPE_CHECKING import for proper typing
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant